import cv2
import bpy
import bpycv
import numpy as np
import os
import json

TEST = True

bpy.data.objects["AM266_001_decoration_A_Octane"]["inst_id"] = 1

scene = bpy.data.scenes["Scene"]

output_data = {}

rootdir = scene.save_path
dataset_name = scene.sof_dataset_name

dataset_dir = os.path.join(rootdir, dataset_name)

os.makedirs(dataset_dir, exist_ok=True)
os.makedirs(os.path.join(dataset_dir, 'mask'), exist_ok=True)

result = bpycv.render_data()
cv2.imwrite(os.path.join(dataset_dir, 'demo.jpg'), result.vis()[..., ::-1])


if TEST:
    pass

else:
    camera_dict = []
    for frame in range(scene.frame_start, scene.frame_end, scene.train_frame_steps):
        
        scene.frame_set(frame)
        filename = os.path.basename( scene.render.frame_path(frame=frame) )
        
        filedir = 'mask'

        frame_data = {
            'file_path': os.path.join(filedir, filename),
        }
        camera_dict.append(frame_data)
        
        
        result = bpycv.render_data()
        cv2.imwrite(os.path.join(rootdir, dataset_name, filedir, filename), np.uint16(result["inst"]))
    
    output_data['frames'] = camera_dict
    json_path = os.path.join(dataset_dir, 'masks_train.json')
    with open(json_path,'w') as file:
        json.dump(output_data, file, indent=4)